"""
Z-Image-Turbo 推理引擎
"""
import torch
import logging
from typing import Optional, Union
from PIL import Image
import io
import base64

logger = logging.getLogger(__name__)


class ZImageTurboInference:
    """Z-Image-Turbo 推理引擎"""
    
    def __init__(self, pipeline, device: str):
        self.pipe = pipeline
        self.device = device
    
    def generate_image(
        self,
        prompt: str,
        height: int = 1024,
        width: int = 1024,
        num_inference_steps: int = 9,
        guidance_scale: float = 0.0,
        seed: Optional[int] = None,
        **kwargs
    ) -> Image.Image:
        """
        生成图像
        
        Args:
            prompt: 文本提示
            height: 图像高度
            width: 图像宽度
            num_inference_steps: 推理步数（Z-Image-Turbo 使用 9 步，实际 8 NFE）
            guidance_scale: Guidance scale（Turbo 模型应使用 0.0）
            seed: 随机种子
            **kwargs: 其他参数
            
        Returns:
            PIL.Image: 生成的图像
        """
        # 设置随机种子
        if seed is not None:
            generator = torch.Generator(self.device).manual_seed(seed)
        else:
            generator = None
        
        # 生成图像
        with torch.no_grad():
            result = self.pipe(
                prompt=prompt,
                height=height,
                width=width,
                num_inference_steps=num_inference_steps,
                guidance_scale=guidance_scale,
                generator=generator,
                **kwargs
            )
        
        return result.images[0]
    
    def generate_image_base64(
        self,
        prompt: str,
        height: int = 1024,
        width: int = 1024,
        num_inference_steps: int = 9,
        guidance_scale: float = 0.0,
        seed: Optional[int] = None,
        image_format: str = "PNG",
        **kwargs
    ) -> str:
        """
        生成图像并返回 base64 编码
        
        Args:
            prompt: 文本提示
            height: 图像高度
            width: 图像宽度
            num_inference_steps: 推理步数
            guidance_scale: Guidance scale
            seed: 随机种子
            image_format: 图像格式 (PNG, JPEG)
            **kwargs: 其他参数
            
        Returns:
            str: Base64 编码的图像
        """
        image = self.generate_image(
            prompt=prompt,
            height=height,
            width=width,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            seed=seed,
            **kwargs
        )
        
        # 转换为 base64
        buffer = io.BytesIO()
        image.save(buffer, format=image_format)
        image_bytes = buffer.getvalue()
        base64_str = base64.b64encode(image_bytes).decode('utf-8')
        
        return base64_str

